//SPDX-FileCopyrightText: Copyright 2022-2024 深圳市同心圆网络有限公司
//SPDX-License-Identifier: GPL-3.0-only

package impl

import (
	"io"
	"net"
	"sync"
	"time"

	"github.com/google/uuid"
	"golang.org/x/crypto/ssh"
)

type LISTEN_TYPE string

const (
	LISTEN_LOCAL  = LISTEN_TYPE("local")
	LISTEN_REMOTE = LISTEN_TYPE("remote")
)

type ForwardClient struct {
	sshClient  *ssh.Client
	li         net.Listener
	ListenType LISTEN_TYPE
	LocalAddr  string
	RemoteAddr string
	connMap    sync.Map
}

func NewForwardClient(sshClient *ssh.Client, localAddr, remoteAddr string, listenType LISTEN_TYPE) (*ForwardClient, error) {
	var li net.Listener
	var err error
	if listenType == LISTEN_LOCAL {
		li, err = net.Listen("tcp", localAddr)
		if err != nil {
			return nil, err
		}
	} else if listenType == LISTEN_REMOTE {
		li, err = sshClient.Listen("tcp", remoteAddr)
		if err != nil {
			return nil, err
		}
	}
	ret := &ForwardClient{
		sshClient:  sshClient,
		li:         li,
		ListenType: listenType,
		LocalAddr:  localAddr,
		RemoteAddr: remoteAddr,
		connMap:    sync.Map{},
	}
	go ret.Run()
	return ret, nil
}

func (fc *ForwardClient) Close() {
	fc.li.Close()
	fc.connMap.Range(func(_, value any) bool {
		connList := value.([]net.Conn)
		for _, conn := range connList {
			conn.Close()
		}
		return true
	})
}

func (fc *ForwardClient) Run() {
	for {
		sConn, err := fc.li.Accept()
		if err != nil {
			return
		}
		var conn net.Conn
		if fc.ListenType == LISTEN_LOCAL {
			conn, err = fc.sshClient.Dial("tcp", fc.RemoteAddr)
			if err != nil {
				continue
			}
		} else if fc.ListenType == LISTEN_REMOTE {
			conn, err = net.DialTimeout("tcp", fc.LocalAddr, 10*time.Second)
			if err != nil {
				continue
			}
		}
		connId := uuid.NewString()
		fc.connMap.Store(connId, []net.Conn{sConn, conn})
		go func(id string) {
			io.Copy(sConn, conn)
			conn.Close()
			sConn.Close()
			fc.connMap.Delete(id)
		}(connId)
		go func(id string) {
			io.Copy(conn, sConn)
			conn.Close()
			sConn.Close()
			fc.connMap.Delete(id)
		}(connId)
	}
}
